{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.parquet as pq\n",
    "import pyarrow as pa\n",
    "import ujson\n",
    "import numpy as np\n",
    "from rich import progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. 处理预训练阶段数据\n",
    "## 1.1 处理Wiki数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "origin_wiki_file = './data/wiki.simple.txt'\n",
    "\n",
    "liness = []\n",
    "with open(origin_wiki_file, 'r', encoding='utf-8') as f:\n",
    "    lines = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['数学：\\n',\n",
       " '数学，是研究数量、结构以及空间等概念及其变化的一门学科，属于形式科学的一种。数学利用抽象化和逻辑推理，从计数、计算、量度、对物体形状及运动的观察发展而成。数学家们拓展这些概念，以公式化新的猜想，以及从选定的公理及定义出发，严谨地推导出一些定理。\\n',\n",
       " '基础数学的知识与运用是生活中不可或缺的一环。对数学基本概念的完善，早在古埃及、美索不达米亚及古印度历史上的古代数学文本便可观见，而在古希腊那里有更为严谨的处理。从那时开始，数学的发展便持续不断地小幅进展，至16世纪的文艺复兴时期，因为新的科学发现和数学革新两者的交互，致使数学的加速发展，直至今日。数学并成为许多国家及地区的教育中的一部分。\\n',\n",
       " '数学在许多领域都有应用，包括科学、工程、医学、经济学和金融学等。数学对这些领域的应用通常被称为应用数学，有时亦会激起新的数学发现，并导致全新学科的发展，例如物理学的实质性发展中建立的某些理论激发数学家对于某些问题的不同角度的思考。数学家也研究纯粹数学，就是数学本身的实质性内容，而不以任何实际应用为目标。许多研究虽然以纯粹数学开始，但其过程中也发现许多可用之处。\\n',\n",
       " '词源。\\n']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lines[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10103707"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "合并词条和内容"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "items, content = [], []\n",
    "key_word, kw_line_idx = '', 0\n",
    "content_start = False  # 词条内容开始标记\n",
    "\n",
    "bos_token, eos_token = '[BOS]', '[EOS]' \n",
    "for i, line in enumerate(lines):\n",
    "    \n",
    "    line_strip = line.strip()\n",
    "\n",
    "    # 词条以冒号`：`结尾\n",
    "    if len(line_strip) > 0 and line_strip[-1] in (':', '：'):\n",
    "        key_word = ''.join(line_strip[: -1])\n",
    "        kw_line_idx = i \n",
    "        continue\n",
    "    \n",
    "    # 词条key_word在下一行，则合并上个词条并保存\n",
    "    if i == kw_line_idx + 1 and key_word in line_strip or i == len(lines) - 1:\n",
    "        txt = ''.join(content)\n",
    "\n",
    "        if len(txt) > 0:\n",
    "            items.append(f\"{txt}{eos_token}\")\n",
    "            \n",
    "        content = []\n",
    "        content.append(f\"{key_word}：\")\n",
    "    \n",
    "    content.append(line)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1008067"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(items)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'地质学：地质学（法语、德语：Geologie；拉丁语、西班牙语：Geologia；源于希腊语，γῆ，和，λoγία）是对地球的起源、历史与结构进行研究的学科。主要研究地球的物质组成、内部构造、外部特征、各圈层间的相互作用和演变历史。在现阶段，由于观察、研究条件的限制，主要以岩石圈为研究对象，并涉及水圈、大气圈、生物圈和岩石圈下更深的部分，以及涉及其他行星和卫星的太空地质学。\\n地球历史学。\\n-zh-hant：zh-hans\\n地质时间尺度涵盖了整个地球历史。其起点开始于最早的太阳系材料生成时，约在4。567Ga，地球的出现约在4。54，Ga，一开始称为冥古宙，最年轻的地质年代则是全新世。\\n地质年代比例。\\n如果把地球诞生到现在的大约45亿年缩小到1年，则人类（Homo，apiens存在的时间只有极短的两分钟。\\n显生宙。\\n地质学学史。\\n对地球的物质成分的研究最早至少可以追溯到古希腊泰奥弗拉斯托斯的著作《论岩石》Peri，Lithon）在古罗马时期，老普林尼详细的描述了常用的一些矿物和金属，还正确的解释了琥珀的来源。\\n一些现代学者（如）认为，现代地质学开始于中世纪伊斯兰世界。比鲁尼就是最早的之一，他的著作有最早的叙述的文章，提出了印度次大陆曾经是海洋的假设。伊斯兰学者伊本·西那对山脉的形成，地震的原因，以及其他一些现代地理学的论题给出了详细的解释，这些内容为日后地质学的发展提供了基础。在中国，博学家沈括（1031–1095）提出陆地形成的假说。他在一个离海洋几百公里远的山中的看到，在一个地质地层里有贝壳类生物化石。由此他推论，陆地是由山脉的侵蚀和淤泥的沉积所形成的。\\n很早以前，地质学的知识比较零星分散。关于这方面的知识，如从地中开采金属、黏土、煤和盐的一些知识，早已为矿工和有关的人们所知晓，而自然哲学家们则大都脱离这些实践，独立形成自己的思辨性的地质理论。\\n地质学在18世纪开始成为一门独立的科学，并在19世纪早期达到成熟阶段。\\n1790年至1830年这一段时期被称为“地质学的英雄时代”在这个时期，在考察岩层顺序以及岩层所含矿物和化石上，人们做了大量工作。在现阶段，由于观察、研究条件的限制，主要以岩石圈为研究对象，并涉及水圈、大气圈、生物圈和岩石圈下更深的部位，以及涉及其他行星和卫星的太空地质学。\\n。工作方法的一大进步表现在用根据化石内容来进行岩层分类。\\n地质学史上有三场著名的争论。\\n其一为水成论与火成论之争，发生在18世纪末。争论的焦点在于岩石的形成理论，一方以德国科学家亚伯拉罕·戈特洛布·维尔纳为代表，强调形成岩石过程中的水的作用；另一方以苏格兰科学家赫屯为代表，强调火的作用。现今已经知道，岩石主要由三大类构成，除了水成为主的沉积岩和火成为主的岩浆岩，还存在一类变质岩。水成过程和火成过程在岩石的形成中都扮演了重要角色。\\n其二为灾变论与渐变论（也称均变论）之争，发生在19世纪早期。持灾变论观点的学者认为，地球历史上曾发生过多次大的灾难，是灾难导致了旧的物种的灭绝和新物种的再创造。持渐变论观点的学者认为，物种演化的动力来自于微弱的地质作用在地球演变过程中的长期积累，不依靠大型的灾难也能够发生。\\n第三场争论是固定论与活动论之争。固定论学说认为地壳的位置是永远不变的，其运动方式以垂直运动为主。传统的地槽、地台学说就是一种固定论学说。活动论学说认为地壳的运动是以水平运动为主的，垂直运动虽然存在，但是是水平运动过程中派生出来的（比如两地块水平挤压处地壳向上隆起）活动论随着板块构造学说的发展，逐渐被学者所认同。\\n[EOS]'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "items[20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将Wiki数据合并为长度固定的行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_txt_cropus_to_chunk_data(texts: list[str], batch_size: int=512 ** 2, max_len: int=320, window_size: int = 2) -> list[str]:\n",
    "    \n",
    "    buffer, buffer_len = [], 0\n",
    "    chunk_data = []\n",
    "\n",
    "    for i, line in enumerate(texts):\n",
    "        buffer_len += len(line)\n",
    "        buffer.append(line)\n",
    "\n",
    "        if buffer_len >= batch_size or i == len(texts) - 1:\n",
    "            buffer_txt = ''.join(buffer)\n",
    "            \n",
    "            # - window_size为滑动窗口，这样每个窗口都包含有window_size个上文\n",
    "            for i in range(0, len(buffer_txt), max_len - window_size):\n",
    "\n",
    "                chunk_data.append(''.join(buffer_txt[i: i + max_len]))\n",
    "            \n",
    "            buffer, buffer_len = [], 0\n",
    "    \n",
    "    return chunk_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2045355\n"
     ]
    }
   ],
   "source": [
    "chunk_data = split_txt_cropus_to_chunk_data(items)\n",
    "print(len(chunk_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['数学：数学，是研究数量、结构以及空间等概念及其变化的一门学科，属于形式科学的一种。数学利用抽象化和逻辑推理，从计数、计算、量度、对物体形状及运动的观察发展而成。数学家们拓展这些概念，以公式化新的猜想，以及从选定的公理及定义出发，严谨地推导出一些定理。\\n基础数学的知识与运用是生活中不可或缺的一环。对数学基本概念的完善，早在古埃及、美索不达米亚及古印度历史上的古代数学文本便可观见，而在古希腊那里有更为严谨的处理。从那时开始，数学的发展便持续不断地小幅进展，至16世纪的文艺复兴时期，因为新的科学发现和数学革新两者的交互，致使数学的加速发展，直至今日。数学并成为许多国家及地区的教育中的一部分。\\n数学在许多领域都有应用，包括科学、工程、医学、',\n",
       " '学、经济学和金融学等。数学对这些领域的应用通常被称为应用数学，有时亦会激起新的数学发现，并导致全新学科的发展，例如物理学的实质性发展中建立的某些理论激发数学家对于某些问题的不同角度的思考。数学家也研究纯粹数学，就是数学本身的实质性内容，而不以任何实际应用为目标。许多研究虽然以纯粹数学开始，但其过程中也发现许多可用之处。\\n词源。\\n西方语言中“数学”一词源自于古希腊语的（其有“学习”学问”科学”还有个较狭义且技术性的意思－\"数学研究\"即使在其语源内。其形容词（意思为\"和学习有关的\"或\"用功的\"亦会被用来指\"数学的\"其在英语中表面上的复数形式，及在法语中的表面复数形式\\'可溯至拉丁文的中性复数\\'由西塞罗译自希腊文复数（此一希腊语被亚里士多德',\n",
       " '多德拿来指\"万物皆数\"的概念。\\n汉字表示的\"数学\"一词大约产生于中国宋元时期。多指象数之学，但有时也含有今天上的数学意义，例如，秦九韶的《数学九章》永乐大典》记，即《数书九章》也被宋代周密所著的《癸辛杂识》记为《数学大略》数学通轨》明代柯尚迁著）数学钥》清代杜知耕著）数学拾遗》清代丁取忠撰）直到1939年，经过中国数学名词审查委员会研究“算学”与“数学”两词的使用状况后，确认以“数学”表示今天意义上的数学含义。\\n历史。\\n数学有着久远的历史。它被认为起源于人类早期的生产活动：中国古代的六艺之一就有\"数\"数学一词在西方有希腊语词源（mathematikós意思是“学问的基础”源于（máthema，科学，知识，学问”\\n史前的人类就已尝试用']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunk_data[0: 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = pa.Table.from_arrays([chunk_data], names=['text'])\n",
    "# compression='GZIP'\n",
    "pq.write_table(table=tb, where='./data/wiki_chunk_320_2.2M.parquet', row_group_size=50000, data_page_size=50000, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wiki_samples = np.random.choice(chunk_data, size=100_0000).tolist()\n",
    "wiki_samples = chunk_data[0: 10_0000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100000"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(wiki_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['台御史杨开鼎的倡议下，于该年十月重建孔庙，于乾隆十六年（1751年）三月完工，而此次工程中大成殿与两庑的规模倍前，格局被记载在隔年王必昌编纂的《重修台湾县志》所附之府学宫图中。然后到了乾隆四十二年（1777年）知府蒋元枢又加以整建，主要是在东大成坊外增建了泮宫坊，大成殿西边增设了府学署，为台南孔庙规模最大的时期，其格局载于《重建台郡各建筑图说》之重建台湾府学图与现存于明伦堂里的台湾府学全图碑上。然而后来孔庙却因为嘉庆与同治年间的地震等因素而逐渐毁损。\\n据日人山田孝使《台南圣庙考》载，在同治元年（1862年）台湾发生大地震，大成殿两庑有倾圯者，修之\"本年元月及五月台湾发生两次地震，特别是五月大地震，使清政府下诏免征田赋。同年官民捐银千',\n",
       " '银千余两。然而比起道光十五年修复工程所费金额的规模，此次规模很小。\\n光绪十三年（1887年）台湾建省。置台湾、台北、台南三府，将原台湾府改称\"台南府\"但直到光绪十四年（1888年）才将\"台湾府学\"改称为\"台南府学\"\\n日治时期。\\n自同治年间至日本领台前，有三十余年台南孔子庙皆无整修，以至于蠹蚀严重。乙末（1895年）改隶后，进入日治时期。因各地抗日义军风起云涌，日军将\"南进讨伐军驻屯所\"设于台南孔子庙内。并在东大成坊设立\"台南舍营病院\"作为军医院所用。直到抗日义军声势较弱后，将台南孔子庙改为\"台湾民政支所职员宿舍\"明治二十九年（1896年）11月4日，在此开办\"台南国语传习所\"\\n明治三十年（1897年）前清举人蔡国琳与商朝凤、许廷光等',\n",
       " '光等人谋议，由台南孔子庙与附属乐局的田园租金收入，共一千二百九十圆来修葺大成殿及其他建物，并修补礼乐器。然而在明治三十一年（1898年）8月，台湾总督府发布\"公学校令\"和\"公学校规则\"在台南孔庙设立\"台南公学校\"提供台湾人小学程度日文教育，明治三十七年（1904年）更名为\"台南第一公学校\"当时校舍的使用，除了大成殿、文昌阁及节孝祠外，其余皆作为学校使用。事实上，大成殿亦被充当教室使用。大正五年（1916年）孔庙计划整修，所以学校于隔年（1917年）8月迁于南大附小树林街现址\\n\\n大正五年（1916年）4月，中华民国福建省指派代表张遵旭来台参观台湾劝业会。将访台18日的见闻撰成《台湾游记》即纪录参观台南孔子庙的境况：现在孔庙改为公学校（',\n",
       " '校（公学校，专为教育台湾人而设）时正教授国语，见黑板上书\"大日本帝国\"尊国体\"美国风\"等字，教员发问，诸生争相对答解讲\"吾等应如何尊重国体\"应如何爱慕国风\"噫！国语之重要如此！然纯为化民根本上之问题，深足畏也（予于1910年曾往朝鲜，参观京城各学校亦皆注意于此）周览大成殿，中置至圣先师之位，景象萧条，盖已废祀久矣。庙旁旧有海东书院，今已倾圮。\\n大正五年（1916年）至六年（1917年）间，日本政府辟建幸町（今南门路）泮宫石坊也移至现今位置（南门路与府中路交叉巷口）\\n大正六年（1917年）7月孔庙整修动工，隔年（1918年）3月整修完毕，由台南厅长枝德二提倡修复，台南仕绅许廷光、杨鹏抟、陈鸿鸣、谢群我等人发起捐献。完工后枝德二撰文立碑',\n",
       " '立碑，采汉文说明整修始末，碑身位于今孔庙东南侧。此次工程主要是将建筑物解体重建，山田孝使的《台南圣庙考》一书载有施工前后的孔庙平面图，其差异主要是在不再复建明伦堂两廊与后方的文昌祠，朱子祠、櫺星门以及礼门义路的围墙也未重建，形成孔庙今天的格局。并由山田孝使撰写《台南圣庙考》将孔子庙沿革、倡修、维护及捐款者等刊印。\\n大正十二年（1923年）日本皇太子裕仁访问台湾，于4月20日访问台南孔子庙。\\n昭和十一年（1936年）左右，孔庙泮池前空地（今忠义国小操场）割给台南神社作为神社外苑。皇民化运动时期，台南孔庙曾被设置日本神棚，被注重\"文化财\"的时任台南市长羽鸟又男撤除，为此曾遭检举为\"非国民\"羽鸟又男任内亦维护恢复了祭孔仪式。\\n昭和十六年（']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wiki_samples[-5:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  1.2 处理百度百科数据库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.parquet as pq\n",
    "import pyarrow as pa\n",
    "import ujson\n",
    "from unicodedata import normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_txt_cropus_to_chunk_data(texts: list[str], batch_size: int=512 ** 2, max_len: int=320, window_size: int = 2) -> list[str]:\n",
    "    \n",
    "    buffer, buffer_len = [], 0\n",
    "    chunk_data = []\n",
    "\n",
    "    for i, line in enumerate(texts):\n",
    "        buffer_len += len(line)\n",
    "        buffer.append(line)\n",
    "\n",
    "        if buffer_len >= batch_size or i == len(texts) - 1:\n",
    "            buffer_txt = ''.join(buffer)\n",
    "            \n",
    "            # - window_size为滑动窗口，这样每个窗口都包含有window_size个上文\n",
    "            for i in range(0, len(buffer_txt), max_len - window_size):\n",
    "\n",
    "                chunk_data.append(''.join(buffer_txt[i: i + max_len]))\n",
    "            \n",
    "            buffer, buffer_len = [], 0\n",
    "    \n",
    "    return chunk_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "bd_baike_563w_file = './data/563w_baidubaike.json'\n",
    "baike_items = []\n",
    "eos_token = '[EOS]' \n",
    "max_len = 320\n",
    "batch_size, batch_cnt = 200_0000, 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'title': '红色食品', 'summary': '红色食品是指食品为红色、橙红色或棕红色的食品。科学家认为,多吃些红色食品可预防感冒。红色食品有红柿椒、西红柿、胡萝卜、红心白薯、红果(山楂)、红苹果、草莓、红枣、老南瓜、红米、柿子等。 有治疗缺铁性贫血和缓解疲劳的作用,对乳腺癌等肿瘤疾病有防治作用,给人以兴奋感,有增加食欲,光洁皮肤,增强表皮细胞再生和防止皮肤衰老,预防感冒等作用。', 'sections': [{'title': '简介', 'content': '红色食品富含番茄红素、胡萝卜素、铁和部分氨基酸,是优质蛋白质、碳水化合物、膳食纤维、B族维生素和多种无机盐的重要来源,可以弥补粳米、白面中的营养缺失。经常食用红色食品,可以进一步提高对主食中营养的利用率,山植等食品还有治疗癌症的功效。被称为“红色生力军。”营养学家认为,红色蔬果最典型的优势在于它们都是富含天然铁质的食物,例如我们常吃的樱桃、大枣等都是贫血患者的天然良药,也适合女性经期失血后的滋补。所以,红色蔬果,女人尽可放心多吃。红色食品中还含有致病微生物的“杀手”——巨噬细胞,可以有效地抵御感冒病毒等微生物,增强人体抵抗感冒的能力。  红色食品\\n在所有的果蔬当中,名声最好的莫过于苹果。西方有“One apple a day,keeps the doctors away.”的说法,因为苹果性情温和,含有各种维生素和微量元素,是所有的水果中最接近完美的一个。\\n还有一种说法:红色食品是相对于绿色食品而言的,指对人体有害的食品,如各种有毒有害、腐败变质、添加非食用物质的食品。红色食品危害人体健康乃至生命安全,对人体健康亮起了红灯,应当大力查处。'}, {'title': '作用', 'content': '这些食品中富含β-胡萝卜素和维生素A,对孩子上皮组织和呼吸道粘膜有很强的保护作用,可提高预防感冒的能力。\\n假如你生来体质较弱,易受感冒病毒的困扰,或者已经被感冒缠上了,红色食品会助你一臂之力,天生具有促进人体健康卫士之一的巨噬细胞活力的功能,巨噬细胞乃是感冒病毒等致病微生物的“杀手”,其活力增强了,感冒病毒自然难以在人体内立足,更谈不上生长繁殖了。至于颜色较辣椒稍浅一些的胡萝卜,所含的胡萝卜素可在体内转化为维生素A,发挥护卫人体上皮组织如呼吸道黏膜的作用,常食之同样可以增强人体抗御感冒的能力。除了红辣椒、胡萝卜外,苋菜、洋葱、红枣、番茄、红薯、山楂、苹果、草莓、老南瓜、红米等亦具此功。'}, {'title': '红色食品与感冒', 'content': '冬令时节,气候寒冷,万物收藏,人的机体生理功能处于降低、抑制、收缩状态,易患感冒,吃红色食品可扶正祛邪,增强免疫力,预防得病。[1]\\n蔬菜中的红萝卜、红辣椒、番茄、红薯、红枣、红苋菜等红色食品中,富含β-胡萝卜素,不但能清除人体氧自由基,而且参与合成维生素A,对人体上皮组织和 呼吸道黏膜有很强的保护作用。推荐以下几个食疗方———\\n★番茄猪肝汤\\n用料:猪肝250克,虾仁25克,蘑菇40克,鸡蛋1只,番茄150克,黄酒、葱段、姜片、胡椒粉、精盐适量。\\n制法:将猪肝切去筋膜洗净,切丁后加上酒、姜汁、蛋液、盐、胡椒粉,搅打成浆。用旺火蒸10—15分钟至结膏。清水加虾仁、黄酒沸煮5分钟后倒入蘑菇、番茄丁和肝膏,再煮沸,调味即可。\\n功用:养肝明目,增强免疫力。用于防感冒、防治夜盲症及免疫力低下者,以及甲亢。\\n方解:猪肝有养肝明目、增强免疫力;蘑菇补益脾胃,益阴养肝,降压,降脂,润燥化痰,增加白细胞;虾仁、番茄均有增强免疫力的食品,番茄可增强营养,减少感冒的发生。\\n★红萝卜炖牛肉\\n用料:牛肉250克、红萝卜250克。\\n制法:牛肉切成小块,加黄酒、姜、葱等配料,再加入红萝卜块,炖熟,即可食用。\\n功用:益气养胃、强健筋骨、增强免疫力。适用于防感冒及免疫力低、虚损消瘦、腰膝酸软者。\\n方解:牛肉补脾胃、益气血、强筋骨,红萝卜增强免疫力,防感冒,健脾,补血,助消化。\\n★蜂蜜红萝卜汁\\n用红萝卜汁与蜂蜜各半制成混合汁剂,每天饮用3次,每次1汤匙。可防治伤风、感冒和咽喉炎。胡萝卜能提供丰富的维生素A,具有促进机体正常生长与繁殖、维护上皮组织、防止感冒及保持视力正常。蜂蜜能补中,润燥,止痛,解毒,清热。'}, {'title': '红色食品与红肉', 'content': '红色食品是指外表呈红色的果蔬和“红肉”类。红色果蔬包括红辣椒、西红柿、红枣、山楂、草莓、苹果等,红色果蔬含有糖和多种维生素,尤其富含维生素C。“红肉”指牛肉、猪肉、羊肉及其制品。\\n红色果蔬中的辣椒具有温中散寒,开胃除湿之功效,辣椒中的辣椒素能刺激唾液和胃肠道消化液的分泌,还能刺激心血管系统,使心跳加快,血液循环加速,因此在寒冷环境有祛风除湿的作用。风寒型感冒病人食用辣椒汤能帮助发汗,有利于感冒的康复,但是胃肠疾病、结核病人则不适合食用。西红柿-在国外享有“金苹果”之称,具有较高的价值。由于西红柿含有94%左右的水分,生吃能防治中暑,止渴生津,凉血解毒的作用,但食西红柿时尽量少放盐,为了避免维生素的破坏,做汤时最好等水开了再下西红柿,而且忌食未成熟的西红柿胃,胃肠虚寒者即慢性腹泻和消化不良者应忌食之。红枣在国外被称为“天然维生素丸”,具有很好的补血功效,能安神和补益脾胃,但胃肠积滞和患有牙齿疾病者应忌食。食用红枣时不宜与鱼同食,同食易引起腹部胀痛。\\n红色食品中的肉类即所谓“红肉”,主要含蛋白质和脂肪及其它无机盐等,因此具有丰富的营养价值。不过“红肉”致癌,世界癌症研究基金会建议食用“红肉”时,每日每人撮入量应少于80克,这是因为“红肉”在烧烤、烙制、煎炸时,其表面产生多种杂环胺——致癌物。'}, {'title': '好处', 'content': '红色不但能让人联想到爱情和激情,还是一种与心脏、大脑和泌尿系统的健康有关的颜色。红色的水果和蔬菜对我们的身体健康大有裨益。 红色的果蔬主要含有丰富的植物化学成分,包括抗细胞因子、抗氧化剂和番茄红素(一种产生红色的色素)。这些成分能够预防癌症,特别是肺癌、前列腺癌和消化道癌。它们还可以延缓衰老,并且有利于防止黄斑变性。黄斑变性是导致65岁以上老年人失明的主要诱因。\\n1、草莓含有80%的水分、丰富的维生素C、少量膳食纤维和钾。由于含糖量很低,因此经常出现在减肥食谱中。大量抗坏血酸、凝集素和果胶使它成为降低血液中胆固醇含量的理想食物。传统医学中草莓可以作为润肤剂、净化剂和利胆剂,还具有镇咳和抗风湿的作用。人们还认为它有抗贫血和滋补的功效,可以促进机体生长。草莓叶子的浸剂还可以用于肠道消炎。\\n2、每天生食6至1 2颗樱桃对痛风和尿酸过多有显著疗效。樱桃的果汁有利于治疗腹泻和结肠炎,所含的抗氧化剂具有保健作用。樱桃含有的主要养分和膳食纤维较少,但是维生素B的含量不低。在矿物质中,钾和铁的含量较高。\\n3、西瓜是水分含量最高的水果,高达其总重量的95%,因此可以促进肾脏更好地发挥功能,将废物和有毒物质排出体外。\\n4、覆盆子含有丰富的维生素E、多种植物营养素和不可溶性纤维。除了具有利尿和通便作用,它还可以用于治疗风湿。\\n5、红苹果富含果胶、糖分和维生素C。此外,由于它是温和的通便剂,所以还具有特殊的医疗效用,可以用于治疗肠道功能紊乱。因此,自然医学认为红苹果可以抗腹泻、贫血和哮喘。它还能够缓解神经系统紧张,促进睡眠。每天晚上吃一个红苹果有助于迅速入睡。它还对希望保持体形的人有用,因为几乎不含脂肪,每100克只有不到58卡路里的热量。据法国图卢兹大学的研究结果显示,每天吃一个大苹果可以在8个星期内使胆固醇水平降低。\\n6、红辣椒中抗氧化剂、硒和维生素C的含量很高,甚至高于柑桔和柠檬等酸味水果。红辣椒所含的膳食纤维能够控制血液中的胆固醇和葡萄糖,还可以改善肠道功能。\\n7、红萝卜有益于治疗呼吸系统疾病,例如咽炎和喉炎,还可以减轻喉咙嘶哑。在柠檬的辅助下,它还可以用来防治哮喘和鼻窦炎。萝卜酒具有清除肾结石和治疗肝脏和胆囊疾患的作用。红萝卜含有钾和少量铁,不含脂肪,每100克只含有15卡路里的热量。它非常适于制作凉拌沙拉,配上柠檬和盐就是一道佳肴。此外,吃红萝卜还可以控制前列腺癌变。\\n8、番茄含有番茄红素和大量抗氧化剂,能够降低患上慢性疾病的危险,尤其是前列腺癌和心血管疾病。番茄具有提神、助消化和抗炎的作用。用它制作的沙拉、酱汁和菜泥可以帮助患有胃炎和胃溃疡的人更好地消化不易消化吸收的食物。番茄的热量很低,含有维生素C。它富含的番茄红素可以防止罹患前列腺癌。如果使用食用油烹调番茄,还可以增强这种功效。\\n'}], 'tags': ['饮食', '食品', '食疗', '科学', '健康', '食品类型'], 'url': 'http://baike.baidu.com/view/0010.htm'}\n"
     ]
    }
   ],
   "source": [
    "with open(bd_baike_563w_file, 'r', encoding='utf-8') as f:\n",
    "    line = f.readline()\n",
    "    line = normalize('NFKC', line)\n",
    "    item = ujson.loads(line)\n",
    "    print(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save to ./data/baike_chunk_320_5.6M_0.parquet\n",
      "save to ./data/baike_chunk_320_5.6M_1.parquet\n",
      "save to ./data/baike_chunk_320_5.6M_2.parquet\n"
     ]
    }
   ],
   "source": [
    "with open(bd_baike_563w_file, 'r', encoding='utf-8') as f:\n",
    "\n",
    "    def process_none(s: str) -> str:\n",
    "        if s: return s\n",
    "        return ''\n",
    "    \n",
    "    while True:\n",
    "        line = f.readline()\n",
    "        if not line: break\n",
    "\n",
    "        item = ujson.loads(line)\n",
    "        cur_txt, cur_len = [], 0\n",
    "\n",
    "        if not item['title']: continue\n",
    "\n",
    "        temp_txt = f\"{item['title']}：{process_none(item['summary'])}\"\n",
    "        \n",
    "        cur_len += len(temp_txt)\n",
    "        cur_txt.append(temp_txt)\n",
    "\n",
    "        for section in item['sections']:\n",
    "\n",
    "            # 太长的截断不要了\n",
    "            if cur_len > max_len:\n",
    "                break\n",
    "            \n",
    "            title = f\"{section['title']}：\" if section['title'] else \"\"\n",
    "            temp_txt = f\"{title}{process_none(section['content'])}\"\n",
    "            \n",
    "            cur_len += len(temp_txt)\n",
    "            cur_txt.append(temp_txt)\n",
    "        \n",
    "        # normalize 处理\\u3000 \\xa0，全角转半角\n",
    "        temp_txt =  normalize('NFKC', ''.join(cur_txt))\n",
    "\n",
    "        if len(temp_txt) > max_len:\n",
    "            # 从 max_len 开始找第一个句号，叹号\n",
    "            n, i = len(temp_txt), max_len\n",
    "            while i < n and temp_txt[i] not in ('。', '！'):\n",
    "                i += 1\n",
    "            temp_txt = ''.join(temp_txt[0: i + 1])\n",
    "\n",
    "        # 添加 eos token\n",
    "        temp_txt = f\"{temp_txt}{eos_token}\"\n",
    "        \n",
    "        baike_items.append( temp_txt )\n",
    "\n",
    "        if len(baike_items) % batch_size == 0:\n",
    "\n",
    "            chunk_data = split_txt_cropus_to_chunk_data(baike_items)\n",
    "            tb = pa.Table.from_arrays([chunk_data], names=['text'])\n",
    "\n",
    "            file_name = f'./data/baike_chunk_320_5.6M_{batch_cnt}.parquet'\n",
    "            pq.write_table(table=tb, where=file_name, row_group_size=50000, )\n",
    "\n",
    "            print(f\"save to {file_name}\")\n",
    "\n",
    "            batch_cnt += 1\n",
    "            baike_items = []\n",
    "\n",
    "    if len(baike_items) > 0:\n",
    "        chunk_data = split_txt_cropus_to_chunk_data(baike_items)\n",
    "        tb = pa.Table.from_arrays([chunk_data], names=['text'])\n",
    "\n",
    "        file_name = f'./data/baike_chunk_320_5.6M_{batch_cnt}.parquet'\n",
    "        pq.write_table(table=tb, where=file_name, row_group_size=50000, )\n",
    "\n",
    "        print(f\"save to {file_name}\")\n",
    "\n",
    "        batch_cnt += 1\n",
    "        baike_items = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bake all lines: 5195460\n"
     ]
    }
   ],
   "source": [
    "file_list = [\n",
    "    f'./data/baike_chunk_320_5.6M_{batch_cnt}.parquet' for batch_cnt in range(3)\n",
    "]\n",
    "\n",
    "line_cnt = 0 \n",
    "for file in file_list:\n",
    "    pf = pq.read_table(file)\n",
    "    line_cnt += pf.num_rows\n",
    "\n",
    "print(f\"bake all lines: {line_cnt}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['连市委办公厅、大连市人民政府办公厅关于印发<旅顺口区人民政府机构改革方案>的通知》(大委办发【2009】29号),设立大连市旅顺口区安全生产监督管理局,正处级建制,为区政府工作部门。[1]主要职责:(一)贯彻落实国家、省、市关于安全生产方面的方针、政策、法律、法规;拟订全区安全生产工作政策和规划;制定发布各行业综合性安全生产规程;指导、协调和监督全区安全生产工作;定期分析和预测全区安全生产形势;负责发布全区安全生产信息;协调解决安全生产中的重大问题。 \\n(二)承担全区安全生产综合监督管理责任,依法行使综合监督管理职权;指导协调、监督检查区政府有关部门和各街道安全生产工作;负责区政府安全生产目标管理考核工作的具体组织实施,监督考核并通',\n",
       " '并通报安全生产控制指标执行情况,监督事故查处和责任追究落实情况。[EOS]陶令秫:táo lìng shú ㄊㄠˊ ㄌㄧㄥˋ ㄕㄨˊ 陶令秫 [EOS]炝炒虾球西兰花:原料::青虾250克、料酒1汤勺、干淀粉2汤勺、橄榄油1汤勺、盐1茶匙、生抽1汤勺。 做法::1、青虾去头去壳去虾腺,从虾背处破刀。 \\n2、清理好的虾仁放在碗中,撒上1/2茶匙底盐腌制入味后,用厨房用纸吸去虾仁身上多余的水分。放入淀粉用手抓几下虾仁,使淀粉上劲包裹住虾仁。上浆后的虾仁静止5-10分钟。 \\n3、煮锅中倒入适量清水烧沸,淋上少许橄榄油,下入西兰花焯烫30秒捞出。 \\n4、放在冷水中浸泡片刻捞出沥干备用。 \\n5、炒锅加热,倒入适量的炒菜油。约三成热的时候下入虾',\n",
       " '入虾仁滑炒,微微变色后捞出备用。 \\n6、炒锅倒入底油,约6成热左右,葱炝锅放入西兰花,大火快手翻炒,放入滑过油的虾仁,顺着锅边烹入一勺料酒,淋上生抽,出锅前再加入1/2勺盐调味后即可起锅装盘。[EOS]刘道辰:刘道辰[1]:男,1964年出生,聊城大学副教授,副院长兼办公室主任,主要从事自然地理学等方面的教学、科研工作;在核心期刊发表科研论文20余篇,参编著作、教材5部,主持参与科研项目10余项。 [EOS]邓慧平:邓慧平[1]:男,1962年出生,理学博士,1996年博士后出站,聊城大学教授;从事气象气候学等方面的教学科研工作,发表科研论文30余篇,2004年10月至2006年11月于美国进行科研合作。 [EOS]李仰东:人物简介',\n",
       " '简介:[1]\\n李仰东,男,中共党员,大专文化,1968年10月1日参加工作。现任赣州市赣江建筑安装工程公司总经理兼党支部书记。李仰东同志从一九六八年组建了赣江修建队,后发展更名为赣江建筑公司,2002年兼并了章贡区南外、水西、湖边、沙石、城东、金源六个建筑公司,组建了赣州市赣江建筑安装集团总公司,先后任队长、经理、党支部书记兼总经理,历任法人代表41年。年产值增到1亿元以上,年纳税600万元以上。企业效益上去了,李仰东不忘报答社会,从1983年开始,赣江公司每年有计划地捐助教育事业,先后为赣州一中建设“夜话亭、东操场”,为赣州三中建大门、建花园,为赣二中修建道路、建宣传栏和花园,共计捐款达33万余元。全家为信丰万隆乡修桥、修村委会、',\n",
       " '会、修路,并建好一座希望小学,共计花费40余万元,为家乡造福。[EOS]枫露皇苑:基本信息:小区均价:17500元/平方米\\n物业类型:别墅 \\n小区地址:通州区北关环岛东北侧1000米\\n物业公司:北京里程物业\\n停车位:0个\\n绿化率:66%\\n周边配套地铁:八通线 通州北苑; \\n公交[1]:667路 855路 316路 388路 649路 675路\\n学校:通州区少年宫,北京物资学院,北京人文大学,北京现代音乐学院\\n医院:通州区潞河医院,通州区中医医院,通州区卫生防疫站,通州区新华医院,中国人民解放军二六三医院本数据来源于百度地图,最终结果以百度地图数据为准。\\n银行:中国工商银行通州区支行、中国银行通州区支行、中国农业银行通州支行 \\n购物:']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunk_data[20: 25]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 处理bell指令数据\n",
    "尝试在预训练阶段加入prompt指令数据，就是尝试在预训练解决加加入部分Sft数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = []\n",
    "eval_data = []\n",
    "eval_size = 1_0000\n",
    "max_len = 400\n",
    "root = 'D:/GitHub/ChatLM-mini-Chinese/data/raw_data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(root + '/bell_open_source/train_3.5M_CN.json', 'r', encoding='utf-8') as f:\n",
    "    for line in f:\n",
    "        item = ujson.loads(line)\n",
    "\n",
    "        if len(item['conversations']) != 2: continue\n",
    "\n",
    "        conversation = item['conversations']\n",
    "        txt = ''\n",
    "        if conversation[0]['from'] =='human':\n",
    "            txt = f\"{conversation[0]['value']}\\n{conversation[1]['value']}\"\n",
    "        else:\n",
    "            txt = f\"{conversation[1]['value']}\\n{conversation[0]['value']}\"\n",
    "        \n",
    "         # 收集测试数据\n",
    "        if len(txt) >= max_len and len(txt) < max_len + 8 and len(eval_data) < eval_size and np.random.rand() <= 0.12:\n",
    "            eval_data.append(txt)\n",
    "            continue\n",
    "            \n",
    "\n",
    "        if len(txt) >= max_len: continue\n",
    "        train_data.append(txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5429 1084177\n"
     ]
    }
   ],
   "source": [
    "print(len(eval_data), len(train_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for file in [root + '/bell_open_source/train_2M_CN.json',  root + '/bell_open_source/Belle_open_source_1M.json']:\n",
    "    with open(file, 'r', encoding='utf-8') as f:\n",
    "        for line in f:\n",
    "            item = ujson.loads(line)\n",
    "\n",
    "            if item['input'].strip() != '':\n",
    "                txt = f\"{item['instruction']}\\n{item['input']}\\n{item['output']}\"\n",
    "            else:\n",
    "                txt = f\"{item['instruction']}\\n{item['output']}\"\n",
    "\n",
    "            # 收集测试数据\n",
    "            if len(txt) >= max_len and len(txt) < max_len + 8 and len(eval_data) < eval_size and np.random.rand() > 0.75:\n",
    "                eval_data.append(txt)\n",
    "                continue\n",
    "            \n",
    "            if len(txt) == 0 or len(txt) >= max_len: continue\n",
    "            train_data.append(\n",
    "                    txt\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000 3150704\n"
     ]
    }
   ],
   "source": [
    "print(len(eval_data), len(train_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['请输入一段描述有关海洋保护的情景对话。\\nPerson A: 哇，这个海滩真美啊！\\nPerson B: 是啊，可惜现在海洋污染越来越严重了。\\nPerson A: 对啊，以前来海边能看到很多海星、海螺等生物，现在好难看到了。\\nPerson B: 这些可爱的生物其实都是海洋生态系统的重要组成部分，如果它们消失了，整个生态系统都会受影响。\\nPerson A: 我们该怎么做才能保护海洋呢？\\nPerson B: 首先，我们要遵守海滩的相关规定，不要随便乱扔垃圾。其次，减少使用一次性塑料制品，比如塑料袋、吸管等，这些东西会进入海洋，危害海洋生物。最后，我们可以积极参与各种海洋保护组织，呼吁更多人加入到保护海洋的行动中来。\\nPerson A: 听起来保护海洋是一项重要的任务，我们一定要做好自己的份内工作。\\nPerson B: 是的，每个人都能为保护海洋贡献自己的一份力量。', '为我寻找5个值得信任的汽车保险公司\\n1. State Farm\\n2. Geico\\n3. Allstate\\n4. Progressive\\n5. Farmers Insurance', 'A公司去年亏损了500万美元，今年净利润增长了50%，今年的净利润是多少？\\n今年净利润为750万美元', '介绍最适合初学者的编程语言\\nPython是最适合初学者学习的编程语言之一。Python具有以下优点：\\n1. 简单易学：Python拥有简单易懂的语法和基本代码结构，可以使初学者更快地掌握编程基础。\\n2. 应用广泛：Python可以用于很多不同的任务和领域，例如数据科学、人工智能、机器学习、网络编程等。\\n3. 强大的社区支持：Python具有庞大的开发者社区，可以获得丰富的资源和支持。\\n4. 多资源学习：Python的学习资源丰富，包括图书、视频、课程等。\\n总之，对于初学者来说，Python是一门非常好的编程语言，它容易入门，应用广泛，有强大的社区支持和丰富的学习资源。', '以下是一道小学数学题： \\n有一条直线，上面有 6 只蚂蚁，其中 4 只向左走，2 只向右走。问：向右走的蚂蚁数量占比是多少？\\nA. 25%\\nB. 33.3%\\nC. 50%\\nD. 66.6%\\n回答该问题。\\n向右走的蚂蚁数量占比为两只蚂蚁中向右走的蚂蚁数量与总蚂蚁数量的比值。\\n先算出总蚂蚁数量：6 只蚂蚁。\\n再算出向右走的蚂蚁数量：2 只蚂蚁。\\n将向右走的蚂蚁数量除以总蚂蚁数量，即得到向右走的蚂蚁数量占比。\\n计算式为：\\n2 ÷ 6 = 0.33\\n所以，向右走的蚂蚁数量占比为 33.3%，选 B。']\n"
     ]
    }
   ],
   "source": [
    "print(train_data[0:5])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = pa.Table.from_arrays([train_data], names=['text'])\n",
    "# compression='GZIP'\n",
    "pq.write_table(table=tb, where=f'./data/bell_pretrain_{max_len}_3M.parquet', row_group_size=20480, data_page_size=20480, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = pa.Table.from_arrays([eval_data], names=['text'])\n",
    "# compression='GZIP'\n",
    "pq.write_table(table=tb, where=f'./data/pretrain_eval_{max_len}_1w.parquet', row_group_size=20480, data_page_size=20480, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. 处理sft阶段数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = []\n",
    "with open('./data/sft_0.8M_CN.json', 'r', encoding='utf-8') as f:\n",
    "    for line in f:\n",
    "        item = ujson.loads(line)\n",
    "\n",
    "        txt = f\"{item['instruction']}{item['output']}\"\n",
    "        \n",
    "        if len(txt) == 0 or len(txt) >= 320: continue\n",
    "        lines.append(\n",
    "                item\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "726475\n"
     ]
    }
   ],
   "source": [
    "print(len(lines))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = pa.Table.from_pylist(lines)\n",
    "# compression='GZIP'\n",
    "pq.write_table(table=tb, where='./data/sft_train_data.parquet', row_group_size=20480, data_page_size=20480, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 统计 token数量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyarrow import parquet as pq\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = PreTrainedTokenizerFast.from_pretrained('./model_save/tokenizer/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 字符数量\n",
    "files = [\n",
    "    './data/baike_chunk_320_5.6M_0.parquet', \n",
    "    './data/baike_chunk_320_5.6M_1.parquet', \n",
    "    './data/baike_chunk_320_5.6M_2.parquet', \n",
    "    './data/bell_pretrain_400_3M.parquet',\n",
    "    # './data/pretrain_eval_400_1w.parquet',\n",
    "]\n",
    "\n",
    "total_char = 0\n",
    "for file in files: \n",
    "    pf = pq.read_table(file)\n",
    "    for row in pf['text']:\n",
    "        total_char += len(row.as_py())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2247061858\n"
     ]
    }
   ],
   "source": [
    "print(total_char)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1837099/1837099 [01:50<00:00, 16567.93it/s]\n",
      "100%|██████████| 1839046/1839046 [01:54<00:00, 16062.48it/s]\n",
      "100%|██████████| 1519315/1519315 [01:28<00:00, 17201.50it/s]\n",
      "100%|██████████| 3150704/3150704 [01:55<00:00, 27381.06it/s]\n"
     ]
    }
   ],
   "source": [
    "total_token = 0\n",
    "buffer = []\n",
    "for file in files: \n",
    "    pf = pq.read_table(file)\n",
    "    n = pf.num_rows\n",
    "    for i, row in tqdm(enumerate(pf['text']), total=n):\n",
    "        buffer.append(row.as_py())\n",
    "\n",
    "        if len(buffer) >= 10000 or i == n - 1:\n",
    "            input_ids = tokenizer(buffer, return_attention_mask=False)['input_ids']\n",
    "            \n",
    "            total_token += sum([len(item) for item in input_ids])\n",
    "            buffer = []\n",
    "\n",
    "if len(buffer) > 0:\n",
    "    input_ids = tokenizer(buffer, return_attention_mask=False)['input_ids']\n",
    "    \n",
    "    total_token += sum([len(item) for item in input_ids])\n",
    "    buffer = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1445852913\n"
     ]
    }
   ],
   "source": [
    "print(total_token)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
